{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 3, Topic 3: DPA on Firmware Implementation of AES"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "NOTE: This lab references some (commercial) training material on [ChipWhisperer.io](https://www.ChipWhisperer.io). You can freely execute and use the lab per the open-source license (including using it in your own courses if you distribute similarly), but you must maintain notice about this source location. Consider joining our training course to enjoy the full experience.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**SUMMARY:** *In the previous lab, you saw how a single bit of information can be used to recover an entire byte of the AES key. Remember, this works due to the S-Box being present in the data flow that we are attacking.*\n",
    "\n",
    "*Next, we'll see how to use power analysis instead of an actual bit value. With this technique, the goal is to separate the traces by a bit in the result of the SBox output (it doesn't matter which one): if that bit is 1, its group of traces should, on average, have higher power consumption during the SBox operation than the other set.*\n",
    "\n",
    "*This is all based on the assumption we discussed in the slides and saw in earlier labs: there is some consistent relationship between the value of bits on the data bus and the power consumption in the device.*\n",
    "\n",
    "**LEARNING OUTCOMES:**\n",
    "\n",
    "* Using a power measurement to 'validate' a possible device model.\n",
    "* Detecting the value of a single bit using power measurement.\n",
    "* Breaking AES using the classic DPA attack.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hold up! Before you continue, check you've done the following tutorials:\n",
    "\n",
    "* ☑ Jupyter Notebook Intro (you should be OK with plotting & running blocks).\n",
    "* ☑ SCA101 Intro (you should have an idea of how to get hardware-specific versions running).\n",
    "* ☑ Breaking AES Using a Single Bit (we'll build on your previous work)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AES Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWHUSKY'\n",
    "CRYPTO_TARGET = 'TINYAES128C'\n",
    "VERSION = 'HARDWARE'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if VERSION == 'HARDWARE':\n",
    "    %run \"Lab 3_3 - DPA on Firmware Implementation of AES (HARDWARE).ipynb\"\n",
    "elif VERSION == 'SIMULATED':\n",
    "    %run \"Lab 3_3 - DPA on Firmware Implementation of AES (SIMULATED).ipynb\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No need to remember the complex model from before - we can instead just jump right into the AES model! Copy your AES model you developed in the previous lab below & run it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ###################\n",
    "# Add your code here\n",
    "# ###################\n",
    "#raise NotImplementedError(\"Add your code here, and delete this.\")\n",
    "\n",
    "# ###################\n",
    "# START SOLUTION\n",
    "# ###################\n",
    "sbox = [\n",
    "    # 0    1    2    3    4    5    6    7    8    9    a    b    c    d    e    f \n",
    "    0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76, # 0\n",
    "    0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0, # 1\n",
    "    0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15, # 2\n",
    "    0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75, # 3\n",
    "    0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84, # 4\n",
    "    0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf, # 5\n",
    "    0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8, # 6\n",
    "    0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2, # 7\n",
    "    0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73, # 8\n",
    "    0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb, # 9\n",
    "    0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79, # a\n",
    "    0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08, # b\n",
    "    0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a, # c\n",
    "    0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e, # d\n",
    "    0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf, # e\n",
    "    0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16  # f\n",
    "]\n",
    "\n",
    "def aes_internal(inputdata, key):\n",
    "    return sbox[inputdata ^ key]\n",
    "# ###################\n",
    "# END SOLUTION\n",
    "# ###################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can verify the model works by running the following blocks, just like last time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Simple test vectors - if you get the check-mark printed all OK.\n",
    "assert(aes_internal(0xAB, 0xEF) == 0x1B)\n",
    "assert(aes_internal(0x22, 0x01) == 0x26)\n",
    "print(\"✔️ OK to continue!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AES Power Watcher"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to send random data to the device, and observe the power consumption during the encryption.\n",
    "\n",
    "The idea is that we will use a capture loop like this:\n",
    "\n",
    "    print(scope)\n",
    "    for i in trange(N, desc='Capturing traces'):\n",
    "        key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "\n",
    "        trace = cw.capture_trace(scope, target, text, key)\n",
    "        if trace is None:\n",
    "            continue\n",
    "        traces.append(trace)\n",
    "        plot.send(trace)\n",
    "\n",
    "    #Convert traces to numpy arrays\n",
    "    trace_array = np.asarray([trace.wave for trace in traces])\n",
    "    textin_array = np.asarray([trace.textin for trace in traces])\n",
    "    known_keys = np.asarray([trace.key for trace in traces])  # for fixed key, these keys are all the same\n",
    "\n",
    "Depending what you are using, you can complete this either by:\n",
    "\n",
    "* Capturing new traces from a physical device.\n",
    "* Reading pre-recorded data from a file.\n",
    "\n",
    "You get to choose your adventure - see the two notebooks with the same name of this, but called `(SIMULATED)` or `(HARDWARE)` to continue. Inside those notebooks you should get some code to copy into the following section, which will define the capture function.\n",
    "\n",
    "Be sure you get the `\"✔️ OK to continue!\"` print once you run the next cell, otherwise things will fail later on!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert(len(trace_array) == 2500)\n",
    "print(\"✔️ OK to continue!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's this data look like? Try plotting a trace or two here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ###################\n",
    "# START SOLUTION\n",
    "# ###################\n",
    "cw.plot(trace_array[0]) * cw.plot(trace_array[1])\n",
    "# ###################\n",
    "# END SOLUTION\n",
    "# ###################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OK interesting - so we've got data! And what about the format of the input data?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(textin_array[0])\n",
    "print(textin_array[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AES Guesser - One Byte"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The attack now needs a way of splitting traces into two groups, depending on the state of a bit in our \"guessed\" value. We're going to start easy by guessing a single byte of the AES key at a time.\n",
    "\n",
    "To start with - define the number of traces & number of points in each trace. You can use the following example code, just run this block:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "numtraces = np.shape(trace_array)[0] #total number of traces\n",
    "numpoints = np.shape(trace_array)[1] #samples per trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you remember from the slides - our algorithm looks like this:    \n",
    "\n",
    "    for key_byte_guess_value in [0, 1, 2, 3, ... 253, 254, 255]:\n",
    "        \n",
    "        one_list = empty list\n",
    "        zero_list = empty list\n",
    "        \n",
    "        for trace_index in [0, 1, 2, 3, ..., numtraces]:\n",
    "        \n",
    "            input_byte = textin_array[trace_index][byte_to_attack]\n",
    "            \n",
    "            #Get a hypothetical leakage list - use aes_internal(guess, input_byte)          \n",
    "\n",
    "            if hypothetical_leakage bit 0 is 1:\n",
    "                append trace_array[trace_index] to one_list\n",
    "            else:\n",
    "                append trace_array[trace_index] to zero_list\n",
    "                \n",
    "        one_avg = average of one_list\n",
    "        zero_avg = average of zero_list\n",
    "\n",
    "        max_diff_value = maximum of ABS(one_avg - zero_avg)\n",
    "        \n",
    "To get the average of your `one_list` and `zero_list` you can use numpy:\n",
    "\n",
    "    import numpy as np\n",
    "    avg_one_list = np.asarray(one_list).mean(axis=0)\n",
    "\n",
    "The important thing here is the `axis=0`, which does an average so the resulting array is done across all traces (not just the average value of one trace, but the average of each point index *across all traces*).\n",
    "\n",
    "To help you do some testing - let me tell you that the correct value of byte 0 is `0x2B`. You can use this to validate that your solution is working on the first byte. If you get stuck - see some hints below (but give it a try first).\n",
    "\n",
    "What you should see is an output of the maximum value between the two average groups be higher for the `0x2B` value. For example, priting the maximum SAD value from an example loop looks like this for me:\n",
    "\n",
    "    Guessing 28: 0.001397\n",
    "    Guessing 29: 0.000927\n",
    "    Guessing 2a: 0.001953\n",
    "    Guessing 2b: 0.005278\n",
    "    Guessing 2c: 0.000919\n",
    "    Guessing 2d: 0.002510\n",
    "    Guessing 2e: 0.001241\n",
    "    Guessing 2f: 0.001242\n",
    "\n",
    "Note the value of `0.005278` for `0x2B` - this is higher than the others which range from `0.000927` to `0.002510`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ###################\n",
    "# Add your code here\n",
    "# ###################\n",
    "#raise NotImplementedError(\"Add Your Code Here\")\n",
    "\n",
    "# ###################\n",
    "# START SOLUTION\n",
    "# ###################\n",
    "import numpy as np\n",
    "mean_diffs = np.zeros(256)\n",
    "\n",
    "guessed_byte = 0\n",
    "\n",
    "for guess in range(0, 256):\n",
    "    \n",
    "    one_list = []\n",
    "    zero_list = []\n",
    "    \n",
    "    for trace_index in range(numtraces):\n",
    "        \n",
    "        #Get a hypothetical leakage list - use aes_internal(guess, input_byte)\n",
    "        hypothetical_leakage = aes_internal(guess, textin_array[trace_index][guessed_byte])\n",
    "    \n",
    "        #Mask off the lowest bit - is it 0 or 1? Depending on that add trace to array\n",
    "        if hypothetical_leakage & 0x01:        \n",
    "            one_list.append(trace_array[trace_index])\n",
    "        else:\n",
    "            zero_list.append(trace_array[trace_index])\n",
    "            \n",
    "    one_avg = np.asarray(one_list).mean(axis=0)\n",
    "    zero_avg = np.asarray(zero_list).mean(axis=0)\n",
    "    mean_diffs[guess] = np.max(abs(one_avg - zero_avg))\n",
    "    \n",
    "    print(\"Guessing %02x: %f\"%(guess, mean_diffs[guess]))\n",
    "    \n",
    "# ###################\n",
    "# END SOLUTION\n",
    "# ###################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hint 1: General Program Flow\n",
    "\n",
    "You can use the following general program flow to help you implement the outer loop above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Hint #1 - General Program Flow\n",
    "import numpy as np\n",
    "mean_diffs = np.zeros(256)\n",
    "\n",
    "guessed_byte = 0\n",
    "\n",
    "for guess in range(0, 256):\n",
    "    \n",
    "    one_list = []\n",
    "    zero_list = []\n",
    "    \n",
    "    for trace_index in range(numtraces):\n",
    "        #Inside here do the steps shown above\n",
    "        pass\n",
    "        \n",
    "    #Do extra steps to average one_list and zero_list        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hint 2: Example of Two Different Key Guesses\n",
    "\n",
    "We aren't fully going to give it away (see `SOLN` notebook if you want that), but here is how you can generate two differences, for `0x2B` and `0xFF`. If you're totally stuck you can use the following code to base what should be inside the loops on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "mean_diffs = np.zeros(256)\n",
    "\n",
    "### Code to do guess of byte 0 set to 0x2B\n",
    "guessed_byte = 0\n",
    "guess = 0x2B\n",
    "   \n",
    "one_list = []\n",
    "zero_list = []\n",
    "    \n",
    "for trace_index in range(numtraces):\n",
    "    hypothetical_leakage = aes_internal(guess, textin_array[trace_index][guessed_byte])\n",
    "\n",
    "    #Mask off the lowest bit - is it 0 or 1? Depending on that add trace to array\n",
    "    if hypothetical_leakage & 0x01:        \n",
    "        one_list.append(trace_array[trace_index])\n",
    "    else:\n",
    "        zero_list.append(trace_array[trace_index])\n",
    "            \n",
    "one_avg = np.asarray(one_list).mean(axis=0)\n",
    "zero_avg = np.asarray(zero_list).mean(axis=0)\n",
    "mean_diffs_2b = np.max(abs(one_avg - zero_avg))\n",
    "\n",
    "print(\"Max SAD for 0x2B: {:1}\".format(mean_diffs_2b))\n",
    "\n",
    "### Code to do guess of byte 0 set to 0xFF\n",
    "guessed_byte = 0\n",
    "guess = 0xFF\n",
    "    \n",
    "one_list = []\n",
    "zero_list = []\n",
    "    \n",
    "for trace_index in range(numtraces):\n",
    "    hypothetical_leakage = aes_internal(guess, textin_array[trace_index][guessed_byte])\n",
    "\n",
    "    #Mask off the lowest bit - is it 0 or 1? Depending on that add trace to array\n",
    "    if hypothetical_leakage & 0x01:        \n",
    "        one_list.append(trace_array[trace_index])\n",
    "    else:\n",
    "        zero_list.append(trace_array[trace_index])\n",
    "            \n",
    "one_avg = np.asarray(one_list).mean(axis=0)\n",
    "zero_avg = np.asarray(zero_list).mean(axis=0)\n",
    "mean_diffs_ff = np.max(abs(one_avg - zero_avg))\n",
    "\n",
    "print(\"Max SAD for 0xFF: {:1}\".format(mean_diffs_ff))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ranking Guesses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll also want to rank some of your guesses (we assume). This will help you identify the most likely value. The best way to do this is build a list of the maximum difference values for each key:\n",
    "\n",
    "    mean_diffs = [0]*256\n",
    "\n",
    "    for key_byte_guess_value in [0, 1, 2, 3, ... 253, 254, 255]:\n",
    "\n",
    "        *** CODE FROM BEFORE***\n",
    "        max_diff_value = maximum of ABS(one_avg - zero_avg)\n",
    "        mean_diffs[key_byte_guess_value] = max_diff_value\n",
    "        \n",
    "If you modify your previous code, it will generate a list of maximum differences in a list. This list will look like:\n",
    "\n",
    "    [0.002921, 0.001923, 0.005131, ..., 0.000984]\n",
    "    \n",
    "Where the *index* of the list is the value of the key guess. We can use `np.argsort` which generates a new list showing the *indicies* that would sort an original list (you should have learned about `argsort` in the previous lab too):\n",
    "\n",
    "So for example, run the following to see it in action on the list `[1.0, 0.2, 3.4, 0.01]`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.argsort([1.0, 0.2, 3.4, 0.01])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This should return `[3, 1, 0, 2`] - that is the order of lowest to highest. To change from highest to lowest, remember you just add `[::-1]` at the end of it like `np.argsort([1.0, 0.2, 3.4, 0.01])[::-1]`.\n",
    "\n",
    "Try using the `np.argsort` function to output the most likely key values from your attack."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Differences"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we move on - you should take a look at various plots of these differences. They will play in something called the *ghost peak* problem.\n",
    "\n",
    "We're going to now define a function called `calculate_diffs()` that implements our attacks (you can replace this with your own function or keep this one for now):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_diffs(guess, byteindex=0, bitnum=0):\n",
    "    \"\"\"Perform a simple DPA on two traces, uses global `textin_array` and `trace_array` \"\"\"\n",
    "    \n",
    "    one_list = []\n",
    "    zero_list = []\n",
    "\n",
    "    for trace_index in range(numtraces):\n",
    "        hypothetical_leakage = aes_internal(guess, textin_array[trace_index][byteindex])\n",
    "\n",
    "        #Mask off the requested bit\n",
    "        if hypothetical_leakage & (1<<bitnum):\n",
    "            one_list.append(trace_array[trace_index])\n",
    "        else:\n",
    "            zero_list.append(trace_array[trace_index])\n",
    "\n",
    "    one_avg = np.asarray(one_list).mean(axis=0)\n",
    "    zero_avg = np.asarray(zero_list).mean(axis=0)\n",
    "    return abs(one_avg - zero_avg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try plotting the difference between various bytes. For byte 0, remember `0x2B` is the correct value. Zoom in on the plots and see how the correct key should have a much larger difference.\n",
    "\n",
    "Sometimes we get *ghost peaks* which are incorrect peaks. So far we're assuming there is a single \"best\" solution for the key - we may need to get fancy and put a threshold whereby we have several candidates for the correct key. For now let's just plot a handful of examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(calculate_diffs(0x2B)) * cw.plot(calculate_diffs(0x2C)) * cw.plot(calculate_diffs(0x2D))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what it should look like:\n",
    "\n",
    "You'll notice when we rank the bytes we just use the maximum value of any peak. There's lots more you could learn from these graphs, such as the location of the peak, or if there are multiple peaks in the graph. But for now we're just going to keep with the "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AES Guesser - All Bytes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alright - good job! You've got a single byte and some DPA plots up. Now let's move onward and guess *all* of the bytes.\n",
    "\n",
    "Doing this requires a little more effort than before. Taking your existing guessing function, you're going to wrap a larger loop around the outside of it like this:\n",
    "\n",
    "    for subkey in range(0,16):\n",
    "        #Rest of code from before!\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import trange\n",
    "import numpy as np\n",
    "\n",
    "#Store your key_guess here, compare to known_key\n",
    "key_guess = []\n",
    "known_key = [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c]\n",
    "\n",
    "for subkey in trange(0, 16, desc=\"Attacking Subkey\"):\n",
    "    # ###################\n",
    "    # Add your code here\n",
    "    # ###################\n",
    "    #raise NotImplementedError(\"Add Your Code Here\")\n",
    "    \n",
    "    # ###################\n",
    "    # START SOLUTION\n",
    "    # ###################\n",
    "    max_diffs = [0]*256\n",
    "    full_diffs = [0]*256\n",
    "    for guess in range(0, 256):\n",
    "        full_diff_trace = calculate_diffs(guess, subkey)\n",
    "        max_diffs[guess] = np.max(full_diff_trace)\n",
    "        full_diffs[guess] = full_diff_trace\n",
    "        \n",
    "    #Get argument sort, as each index is the actual key guess.\n",
    "    sorted_args = np.argsort(max_diffs)[::-1]\n",
    "    \n",
    "    #Keep most likely\n",
    "    key_guess.append(sorted_args[0])\n",
    "    \n",
    "    #Print results\n",
    "    print(\"Subkey %2d - most likely %02X (actual %02X)\"%(subkey, key_guess[subkey], known_key[subkey]))\n",
    "    \n",
    "    #Print other top guesses\n",
    "    print(\" Top 5 guesses: \")\n",
    "    for i in range(0, 5):\n",
    "        g = sorted_args[i]\n",
    "        print(\"   %02X - Diff = %f\"%(g, max_diffs[g]))\n",
    "    \n",
    "    print(\"\\n\")\n",
    "    \n",
    "    # ###################\n",
    "    # END SOLUTION\n",
    "    # ###################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🥳🥳🥳🥳🥳🥳🥳🥳🥳🥳🥳🥳🥳\n",
    "Congrats - you did it!!!!\n",
    "\n",
    "Hopefully the above worked - but we're going to go a little further to understand how to apply this in case it didn't work right away (or it almost worked)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ghost Peaks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Maybe the previous didn't actually recover the full key? No need to worry - there are a few reasons for this. One artifact of a DPA attack is you get another strong peak that isn't the correct key (which can be a ghost peak).\n",
    "\n",
    "We're going to get into more efficient attacks later, but for now, let's look at some solutions:\n",
    "\n",
    "* Increase the number of traces recorded.\n",
    "* Change the targetted bit (& combine solutions from multiple bits).\n",
    "* Window the input data.\n",
    "\n",
    "The first one is the brute-force option: go from 2500 to 5000 or even 10000 power traces. As you add more data, you may find the problem is reduced. But real ghost peaks may not disappear, so we need to move onto other solutions.\n",
    "\n",
    "Before we begin - we're going to give you a \"known good\" DPA attack script we're going to build on. This uses the `calculate_diffs()` function defined earlier.\n",
    "\n",
    "Run the following block (will take a bit of time):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import trange\n",
    "import numpy as np\n",
    "\n",
    "#Store your key_guess here, compare to known_key\n",
    "key_guess = []\n",
    "known_key = [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c]\n",
    "\n",
    "#Which bit to target\n",
    "bitnum = 0\n",
    "\n",
    "full_diffs_list = []\n",
    "\n",
    "for subkey in trange(0, 16, desc=\"Attacking Subkey\"):\n",
    "    \n",
    "    max_diffs = [0]*256\n",
    "    full_diffs = [0]*256\n",
    "\n",
    "    for guess in range(0, 256):\n",
    "        full_diff_trace = calculate_diffs(guess, subkey, bitnum)\n",
    "        max_diffs[guess] = np.max(full_diff_trace)\n",
    "        full_diffs[guess] = full_diff_trace\n",
    "    \n",
    "    #Make copy of the list\n",
    "    full_diffs_list.append(full_diffs[:])\n",
    "    \n",
    "    #Get argument sort, as each index is the actual key guess.\n",
    "    sorted_args = np.argsort(max_diffs)[::-1]\n",
    "    \n",
    "    #Keep most likely\n",
    "    key_guess.append(sorted_args[0])\n",
    "    \n",
    "    #Print results\n",
    "    print(\"Subkey %2d - most likely %02X (actual %02X)\"%(subkey, sorted_args[0], known_key[subkey]))\n",
    "    \n",
    "    #Print other top guesses\n",
    "    print(\" Top 5 guesses: \")\n",
    "    for i in range(0, 5):\n",
    "        g = sorted_args[i]\n",
    "        print(\"   %02X - Diff = %f\"%(g, max_diffs[g]))\n",
    "    \n",
    "    print(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This block should now print some *next top guesses* - in this case just the next top 5 guesses, but you can extend this if you wish. It's also keeping a copy of all the *difference* traces (unlike before where it threw them away)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting Peaks\n",
    "\n",
    "After it runs, select a subkey that is either wrong or has very close \"next best guesses\". For example, the following shows the output for Subkey 5 is actually wrong - the correct guess (`0xAE`) has been ranked as option 5.\n",
    "\n",
    "    Subkey  5 - most likely CB (actual AE)\n",
    "     Top 5 guesses: \n",
    "       CB - Diff = 0.003006\n",
    "       C5 - Diff = 0.002984\n",
    "       AE - Diff = 0.002739\n",
    "       3C - Diff = 0.002674\n",
    "       2F - Diff = 0.002511\n",
    "\n",
    "You can find the full diff in the `full_diffs_list` array. If you index this array it will give you every guess for a given subkey (for example `full_diffs_list[5]` is the 5th subkey guess outputs).\n",
    "\n",
    "Using `full_diffs_list[N]` to get your selected subkey, plot the correct key by plotting `full_diffs_list[N][0xCORRECT]` in green as the *last* (so it appears on top). Plot a few other highly ranked guesses before that. In my example, this would look like:\n",
    "\n",
    "    %matplotlib notebook\n",
    "    import matplotlib.pylab as plt\n",
    "\n",
    "    plt.plot(full_diffs_list[5][0xC5], 'r')\n",
    "    plt.plot(full_diffs_list[5][0xCB], 'r')\n",
    "    plt.plot(full_diffs_list[5][0xAE], 'g')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.plot(full_diffs_list[5][0xC5]) * cw.plot(full_diffs_list[5][0xCB]) * cw.plot(full_diffs_list[5][0xAE])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Zoom in on the window, and you should notice there is a location where the correct peak is *higher* than the incorrect peaks. If you want to plot all the traces (this will get slow!) for a given trace, we can do so as the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = cw.plot()\n",
    "subkey = 12\n",
    "for guess in range(0, 256):\n",
    "    fig *= cw.plot(full_diffs_list[subkey][guess])\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on your hardware, the previous may show a single nice large spike, or multiple large spikes. If we have the ghost peak problem you've probably got multiple spikes. The incorrect peaks may trail behind the correct locations -- we can first plot the correct locations by looking at the known key. The following will do that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = cw.plot()\n",
    "for subkey in range(0, 16):\n",
    "    fig *= cw.plot(full_diffs_list[subkey][known_key[subkey]])\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Windowing Peaks\n",
    "\n",
    "The final trick here - see if there is some way to \"window\" the data that could be useful. For example, looking at the peaks you might notice that the correct peaks are always coming at 60 cycle offsets, with the first peak around sample 1100 (these will be different for your hardware).\n",
    "\n",
    "So we could modify the loop to only look at differences after this point:\n",
    "\n",
    "    for guess in range(0, 256):\n",
    "        full_diff_trace = calculate_diffs(guess, subkey, bitnum)\n",
    "        full_diff_trace = full_diff_trace[(1010 + subkey*60):]\n",
    "        max_diffs[guess] = np.max(full_diff_trace)\n",
    "        full_diffs[guess] = full_diff_trace\n",
    "        \n",
    "Copy the full DPA attack here - and try it out! See if you can get the correct key to come out for every byte."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ###################\n",
    "# Add your code here\n",
    "# ###################\n",
    "#raise NotImplementedError(\"Add Your Code Here\")\n",
    "    \n",
    "# ###################\n",
    "# START SOLUTION\n",
    "# ###################\n",
    "from tqdm.notebook import trange\n",
    "import numpy as np\n",
    "\n",
    "#Store your key_guess here, compare to known_key\n",
    "key_guess = []\n",
    "known_key = [0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c]\n",
    "\n",
    "#Which bit to target\n",
    "bitnum = 0\n",
    "\n",
    "full_diffs_list = []\n",
    "\n",
    "for subkey in trange(0, 16, desc=\"Attacking Subkey\"):\n",
    "    \n",
    "    max_diffs = [0]*256\n",
    "    full_diffs = [0]*256\n",
    "\n",
    "    for guess in range(0, 256):\n",
    "        full_diff_trace = calculate_diffs(guess, subkey, bitnum)\n",
    "        full_diff_trace = full_diff_trace[(0 + subkey*0):]\n",
    "        max_diffs[guess] = np.max(full_diff_trace)\n",
    "        full_diffs[guess] = full_diff_trace\n",
    "    \n",
    "    #Make copy of the list\n",
    "    full_diffs_list.append(full_diffs[:])\n",
    "    \n",
    "    #Get argument sort, as each index is the actual key guess.\n",
    "    sorted_args = np.argsort(max_diffs)[::-1]\n",
    "    \n",
    "    #Keep most likely\n",
    "    key_guess.append(sorted_args[0])\n",
    "    \n",
    "    #Print results\n",
    "    print(\"Subkey %2d - most likely %02X (actual %02X)\"%(subkey, key_guess[subkey], known_key[subkey]))\n",
    "    \n",
    "    #Print other top guesses\n",
    "    print(\" Top 5 guesses: \")\n",
    "    for i in range(0, 5):\n",
    "        g = sorted_args[i]\n",
    "        print(\"   %02X - Diff = %f\"%(g, max_diffs[g]))\n",
    "    \n",
    "    print(\"\\n\")\n",
    "    \n",
    "# ###################\n",
    "# END SOLUTION\n",
    "# ###################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attacking Other Bits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far we only looked at bit 0 $-$ but there are more bits involved here! You can first just try another bit that might be present, maybe they simply work better?\n",
    "\n",
    "But you can also combine multiple bits by creating a most likely solution that applies across *all* bits.\n",
    "\n",
    "The first one is easy to try out, as we defined the bit to attack in the previous script"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second option is a little more advanced. You can give it a try $-$ but in practice, if you are trying to combine multiple bits, a more effective method called the CPA attack will be used."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusions & Next Steps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You've now seen how a DPA attack be be performed using a basic Python script. We'll experience much more effective attacks once we look at the CPA attack.\n",
    "\n",
    "If you want to perform these attacks in practice, the Python code here isn't the most efficient! We'll look at faster options in later courses."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "<small>NO-FUN DISCLAIMER: This material is Copyright (C) NewAE Technology Inc., 2015-2020. ChipWhisperer is a trademark of NewAE Technology Inc., claimed in all jurisdictions, and registered in at least the United States of America, European Union, and Peoples Republic of China.\n",
    "\n",
    "Tutorials derived from our open-source work must be released under the associated open-source license, and notice of the source must be *clearly displayed*. Only original copyright holders may license or authorize other distribution - while NewAE Technology Inc. holds the copyright for many tutorials, the github repository includes community contributions which we cannot license under special terms and **must** be maintained as an open-source release. Please contact us for special permissions (where possible).\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.</small>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if PLATFORM != \"CWLITEXMEGA\":\n",
    "    assert key_guess == known_key, \"Failed to break key, expected vs got:\\n{}\\n{}\".format(known_key, key_guess)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
